import onnx
from onnx import helper
import os.path as osp


if __name__ == '__main__':
    decoder_onnx_path = '/weights/decoder.onnx'
    decoder_onnx = onnx.load(decoder_onnx_path)
    graph = decoder_onnx.graph
    for node in graph.node:
        if node.op_type == 'Tile':
            # 获取 Tile 节点的输入和输出，
            # Tile 节点​​的作用是​​沿指定维度复制输入张量，它的功能类似于 NumPy 的 np.tile() 或 PyTorch 的 torch.repeat()，用于扩展张量的形状和数据
            input_name = node.input[0]
            repeats_name = node.input[1]
            output_name = node.output[0]

            # 创建新的op来替代 Tile
            concat_node = helper.make_node(
                'Concat',
                inputs=[input_name], # 这里试验，发现并不需要repeats
                outputs=[output_name],
                axis=0
            )

            # 替换 Tile 节点
            graph.node.remove(node)
            graph.node.append(concat_node)

    onnx.save(decoder_onnx, osp.join(osp.dirname(decoder_onnx_path), 'decoder_tile2concat.onnx'))
